#include <linux/init.h>
#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/kprobes.h>
#include <linux/uio.h>
#include <linux/blk_types.h>
#include <linux/blkdev.h>
#include <linux/blk-mq.h>

MODULE_AUTHOR("mykprobe");
MODULE_LICENSE("GPL");
MODULE_DESCRIPTION("My kprobe");

/**
 *  /proc/kallsyms
 */
#define TRACE_FUNC "blk_mq_bio_to_request"

int kprobe_pre_handler(struct kprobe *kp, struct pt_regs *regs)
{
    struct request *rq = NULL;
    struct gendisk *disk = NULL;
    struct request_queue *q = NULL;
    unsigned long args1 = 0;
    // unsigned long args2 = 0;
    // unsigned long args3 = 0;
    // unsigned long args4 = 0;
    // unsigned long args5 = 0;
    // unsigned long args6 = 0;
    // unsigned long args7 = 0;

    args1 = regs->di;
    // args2 = regs->si;
    // args3 = regs->dx;
    // args4 = regs->cx;
    // args5 = regs->ax;
    // args6 = regs->r8;
    // args7 = regs->r9;
    rq = (struct request *)args1;
    if (rq == NULL) {
        return 0;
    }
    q = rq->q;
    disk = q->disk;
    if (disk == NULL) {
        return 0;
    }
    if (strncmp(disk->disk_name, "sdc", 3) != 0) {
        return 0;
    }

    printk("rq: %p, start_time_ns: %llu, tag: %d, disk: %s, q: %p\n", 
            rq, rq->start_time_ns, rq->tag, disk->disk_name, q);
    return 0;
}

void kprobe_post_handler(struct kprobe *kp, struct pt_regs *regs, unsigned long flags)
{
    // printk("kprobe_post_handler\n");
    return;
}

int kretprobe_handler(struct kretprobe_instance *krp, struct pt_regs *regs)
{
	unsigned long ret_prt = regs_return_value(regs);
    // printk("blk_mq_get_tag return: %u\n", (unsigned int)ret_prt);
    return 0;
}

static struct kprobe kp = {
    .symbol_name = TRACE_FUNC,
    .pre_handler = kprobe_pre_handler,
    .post_handler = kprobe_post_handler
};

struct kretprobe krp = {
    .kp = {
        .symbol_name = TRACE_FUNC,
    },
    .handler = kretprobe_handler
};

static int __init my_kprobe_init(void)
{
    int ret = 0;
    ret = register_kprobe(&kp);
    ret = register_kretprobe(&krp);
    printk(KERN_INFO "kprobe: start probe.\n");
    return 0;
}

static void __exit my_kprobe_exit(void)
{
    unregister_kprobe(&kp);
    unregister_kretprobe(&krp);
    printk(KERN_INFO "kprobe: end probe.\n");
}

module_init(my_kprobe_init)
module_exit(my_kprobe_exit)
